package sshex

import (
	"fmt"
	"os"
	"sync"

	"gitee.com/shanyou/sshplus/rlog"
	"gitee.com/shanyou/sshplus/utils"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
)

type CertPublicKeyAuth struct {
	Logger     logrus.FieldLogger
	Pubs       map[string][]ssh.PublicKey
	mu         sync.Mutex
	TFAuth     TwoFactorAuth
	TOTPEnable bool
}

func NewCertPublicKeyAuth() *CertPublicKeyAuth {
	l := rlog.Logger.WithField("com", "cert-public-key-auth")
	pubs := make(map[string][]ssh.PublicKey)
	cp := &CertPublicKeyAuth{
		Logger: l,
		Pubs:   pubs,
	}
	return cp
}

func (cp *CertPublicKeyAuth) Config(sshServer *SSHServer) error {
	sshServer.PublicKeyHandler = cp.PublicKeyHandler
	sshServer.KeyboardInteractiveHandler = cp.KeyboardInteractiveHandler
	return nil
}

func (cp *CertPublicKeyAuth) AddCertBytes(keybytes []byte) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	pubs, err := utils.LoadSSHPubKeyBytes(keybytes)
	if err != nil {
		return err
	}

	for n, p := range pubs {
		cp.Pubs[n] = p
	}
	return nil
}

func (cp *CertPublicKeyAuth) AddCertFile(keyPath string) error {
	cp.mu.Lock()
	defer cp.mu.Unlock()
	authorizedKeysBytes, err := os.ReadFile(keyPath)
	if err != nil {
		return err
	}
	return cp.AddCertBytes(authorizedKeysBytes)
}

func (cp *CertPublicKeyAuth) KeyboardInteractiveHandler(conn ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	if !cp.TOTPEnable {
		return nil, ErrPermissionDenied
	}

	user := conn.User()
	ans, err := challenge(user,
		"Topt code verify",
		[]string{"Please input totp code:\n"},
		[]bool{true})
	if err != nil {
		return nil, fmt.Errorf("ssh totp err %v", err)
	}
	valid := cp.TFAuth.Validate(user, ans[0])
	if valid {
		return nil, nil
	}

	return nil, ErrPermissionDenied
}

func (cp *CertPublicKeyAuth) PublicKeyHandler(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	fingerprint := fmt.Sprintf("%v %v", key.Type(), ssh.FingerprintSHA256(key))
	user := conn.User()
	cid := fmt.Sprintf("conn(%v->%v, user=%v, key=%v)", conn.RemoteAddr(), conn.LocalAddr(), user, fingerprint)
	cp.Logger.Infof("%v auth attempt", cid)
	cert, ok := key.(*ssh.Certificate)
	if !ok {
		cp.Logger.Debugf("auth attempt, unsupported key type")
		return nil, fmt.Errorf("unsupported key type")
	}
	if cert.KeyId != conn.User() {
		cp.Logger.Debugf("do not support login name no equals to connection user")
		return nil, fmt.Errorf("connection user is not keyid")
	}
	checker := ssh.CertChecker{
		IsUserAuthority: func(auth ssh.PublicKey) bool {
			for _, pubs := range cp.Pubs {
				for _, pub := range pubs {
					if utils.KeysEqual(pub, auth) {
						return true
					}
				}
			}
			return false
		},
	}
	perm, err := checker.Authenticate(conn, key)
	if err != nil {
		return perm, err
	}

	if cp.TOTPEnable {
		// need topt verify
		err = &ssh.PartialSuccessError{
			KeyboardInteractiveCallback: cp.KeyboardInteractiveHandler,
		}
	}
	return perm, err
}
